import cv2
import os
import io
import json
import config
import numpy as np
import paddle
from paddleocr import PaddleOCR
from PIL import Image, ImageDraw, ImageFont

configs = config.configs
# 定义自定义模型下载路径
ocr_config = configs.get('ocr')
ocr_path = configs.get('ocr_path')
custom_model_dir = ocr_path.get('module')
use_gpu = ocr_config.get('use_gpu')
if use_gpu:
    paddle.set_device('gpu')
# 初始化PaddleOCR，指定使用CPU，并设置模型下载路径
ocr = PaddleOCR(
    use_gpu=ocr_config.get('use_gpu'),
    use_angle_cls=ocr_config.get('use_angle_cls'),
    lang=ocr_config.get('lang'),
    det_model_dir=os.path.join(custom_model_dir, 'det'),  # 检测模型
    rec_model_dir=os.path.join(custom_model_dir, 'rec'),  # 识别模型路径
    cls_model_dir=os.path.join(custom_model_dir, 'cls')  # 分类模型路径
)


def ocr_by_path(file_path):
    # 使用open()函数以二进制模式打开文件
    with io.open(file_path, mode="rb") as file:
        # file对象现在代表了文件的输入流
        file_stream = file
        return ocr_file_stream(file_stream)


def ocr_file_stream(file_stream):
    # 使用PIL的Image.open()直接从文件流中读取图像
    image_pil = Image.open(io.BytesIO(file_stream.read()))
    # 将PIL图像转换为OpenCV格式（BGR通道顺序）
    image_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    # 获取识别结果
    results = ocr.ocr(image_cv)
    # 创建一个ImageDraw对象用于在图像上绘制
    draw = ImageDraw.Draw(image_pil)
    # 设置字体文件，确保包含中文字符支持
    font_file = ocr_path.get('font')  # 替换为您的中文字体文件路径
    font_size = 40  # 根据实际情况调整字体大小
    font = ImageFont.truetype(font_file, font_size)
    # 初始化一个列表来存储识别结果
    result_list = []
    # 输出识别的文字及其位置坐标，并在图像上绘制矩形框和显示文字
    for i, content in enumerate(results):
        if content is None:
            continue
        for position, line in content:
            text, score = line  # 解析识别结果，包括文字内容和置信度
            print(f"Position: {position}")
            print(f"Text: {text} (Confidence: {score})")
            # 将单条识别结果封装为一个字典
            single_result = {
                "position": position,
                "text": text,
                "confidence": score
            }
            result_list.append(single_result)

            # 将 Position 转换为 draw.rectangle() 需要的左上角和右下角坐标
            pt1 = (int(position[0][0]), int(position[0][1]))  # 左上角 (x1, y1)
            pt2 = (int(position[2][0]), int(position[2][1]))  # 右下角 (x2, y2)
            xy = [pt1, pt2]
            # 绘制矩形框
            draw.rectangle((position[0][0], position[0][1], position[2][0], position[2][1]), outline=(0, 255, 0),
                           width=2)

            # 显示文字
            text_padding = 48  # 文字与矩形框的距离

            # 计算文字左上角的坐标，假设我们将其放置在矩形框的上方中央
            text_top_left = (int((pt1[0] + pt2[0]) / 3), pt1[1] - text_padding)

            # 使用 ImageDraw 的 text 方法绘制中文文本
            draw.text(text_top_left, text, fill=(0, 0, 255), font=font)
    # 将结果列表序列化为 JSON 字符串
    json_string = json.dumps(result_list, ensure_ascii=False)

    # 输出 JSON 字符串
    print(json_string)
    # 将PIL图像转换回OpenCV格式以供显示
    img_cv_with_text = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)

    # 先创建窗口，再获取其尺寸
    cv2.namedWindow("Text Detection and Recognition")
    screen_width, screen_height = cv2.getWindowImageRect("Text Detection and Recognition")[2:]

    # 调整图像显示窗口大小以适应屏幕
    img_resized = cv2.resize(img_cv_with_text, (screen_width, screen_height), interpolation=cv2.INTER_AREA)

    # 显示带有识别结果的图像
    cv2.imshow("Text Detection and Recognition", img_resized)
    # cv2.waitKey(0)  # 等待用户按键关闭窗口
    cv2.destroyAllWindows()

    return result_list
